import asyncio
from mcp.client.stdio import stdio_client
from mcp import ClientSession, StdioServerParameters
from hello_agents import SimpleAgent, HelloAgentsLLM, Message
from typing import Optional

class MCPAgent(SimpleAgent):
    """
    一个可以调用MCP服务器工具的Agent
    """
    
    def __init__(
        self,
        name: str,
        llm: HelloAgentsLLM,
        system_prompt: Optional[str] = None,
        mcp_server_path: str = "./mcp_server.py"
    ):
        super().__init__(name, llm, system_prompt)
        self.mcp_server_path = mcp_server_path
        self.server_params = StdioServerParameters(
            command="python",
            args=[mcp_server_path],
            env=None
        )
        
    async def _call_mcp_tool(self, tool_name: str, **kwargs):
        """
        调用MCP服务器上的工具
        """
        try:
            async with stdio_client(self.server_params) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    result = await session.call_tool(tool_name, kwargs)
                    return result.content
        except Exception as e:
            return f"调用MCP工具时出错: {str(e)}"
    
    
    
    def run_with_mcp_tools(self, input_text: str, **kwargs) -> str:
        """
        运行Agent并调用MCP工具
        """
        print(f"🤖 {self.name} 正在处理: {input_text}")
        
        # 构建系统提示词，告知LLM可以使用的MCP工具
        system_prompt = self.system_prompt or "你是一个有用的AI助手。"
        mcp_tools_info = """
        
你可以使用以下MCP工具：
1. calculate - 计算数学表达式，参数: expression (数学表达式)
2. get_current_time - 获取当前时间，无需参数
3. auto_generate_word - 生成单词，参数: prompt (提示词)

当需要使用工具时，请在回复中明确说明你将要调用哪个工具以及原因。
        """
        
        enhanced_system_prompt = system_prompt + mcp_tools_info
        
        # 构建消息历史
        messages = [{"role": "system", "content": enhanced_system_prompt}]
        
        for msg in self._history:
            messages.append({"role": msg.role, "content": msg.content})
            
        messages.append({"role": "user", "content": input_text})
        
        # 获取LLM响应
        response = self.llm.invoke(messages, **kwargs)
        
        # 保存对话历史
        self.add_message(Message(input_text, "user"))
        self.add_message(Message(response, "assistant"))
        print(f"✅ {self.name} 响应完成")
        
        return response
    
    async def calculate_expression(self, expression: str) -> str:
        """
        调用MCP服务器的计算工具
        """
        return await self._call_mcp_tool("calculate", expression=expression)
    
    async def get_current_time_from_mcp(self) -> str:
        """
        调用MCP服务器获取当前时间
        """
        return await self._call_mcp_tool("get_current_time")
    
    async def generate_word(self, prompt: str) -> str:
        """
        调用MCP服务器生成单词
        """
        return await self._call_mcp_tool("auto_generate_word", prompt=prompt)

# 同步包装器函数
def call_mcp_tool_sync(agent: MCPAgent, tool_name: str, **kwargs):
    """
    同步方式调用MCP工具
    """
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        result = loop.run_until_complete(agent._call_mcp_tool(tool_name, **kwargs))
        return result
    finally:
        loop.close()

if __name__ == "__main__":
    # 示例用法
    from dotenv import load_dotenv
    load_dotenv()
    
    # 创建LLM实例
    llm = HelloAgentsLLM()
    
    # 创建MCP Agent
    mcp_agent = MCPAgent(
        name="MCP助手",
        llm=llm,
        system_prompt="你是一个可以调用各种工具的智能助手。"
    )
    
    # 示例：直接调用MCP工具
    print("=== 直接调用MCP工具示例 ===")
    
    # 计算数学表达式
    calc_result = call_mcp_tool_sync(mcp_agent, "calculate", expression="15 * 8 + 32")
    print(f"计算结果: {calc_result}")
    
    # 获取当前时间
    time_result = call_mcp_tool_sync(mcp_agent, "get_current_time")
    print(f"当前时间: {time_result}")
    
    # 生成单词
    word_result = call_mcp_tool_sync(mcp_agent, "auto_generate_word", prompt="科技相关词汇")
    print(f"生成单词: {word_result}")
    
    print("\n=== 使用Agent处理自然语言输入 ===")
    # 使用Agent处理自然语言输入
    response = mcp_agent.run_with_mcp_tools("请计算一下125乘以8再减去50的结果是多少？")
    print(f"Agent响应: {response}")